"""
可视化工具
训练过程和结果的可视化
"""
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pathlib import Path
import seaborn as sns

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class Visualizer:
    """训练可视化器"""
    
    def __init__(self, save_dir):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        
        # 设置样式
        sns.set_style("whitegrid")
        plt.rcParams['figure.figsize'] = (12, 8)
    
    def plot_training_history(self, history, save_name='training_history.png'):
        """
        绘制训练历史
        
        参数:
            history: 训练历史字典
            save_name: 保存文件名
        """
        fig, axes = plt.subplots(2, 3, figsize=(18, 10))
        fig.suptitle('Training History', fontsize=16, fontweight='bold')
        
        epochs = range(1, len(history['train_loss']) + 1)
        
        # 1. Loss曲线
        axes[0, 0].plot(epochs, history['train_loss'], 'b-', label='Train Loss', linewidth=2)
        axes[0, 0].plot(epochs, history['val_loss'], 'r-', label='Val Loss', linewidth=2)
        axes[0, 0].set_xlabel('Epoch')
        axes[0, 0].set_ylabel('Loss')
        axes[0, 0].set_title('Loss Curve')
        axes[0, 0].legend()
        axes[0, 0].grid(True)
        
        # 2. MAE曲线
        train_mae = [m['MAE'] for m in history['train_metrics']]
        val_mae = [m['MAE'] for m in history['val_metrics']]
        axes[0, 1].plot(epochs, train_mae, 'b-', label='Train MAE', linewidth=2)
        axes[0, 1].plot(epochs, val_mae, 'r-', label='Val MAE', linewidth=2)
        axes[0, 1].set_xlabel('Epoch')
        axes[0, 1].set_ylabel('MAE')
        axes[0, 1].set_title('Mean Absolute Error')
        axes[0, 1].legend()
        axes[0, 1].grid(True)
        
        # 3. R²曲线
        train_r2 = [m['R²'] for m in history['train_metrics']]
        val_r2 = [m['R²'] for m in history['val_metrics']]
        axes[0, 2].plot(epochs, train_r2, 'b-', label='Train R²', linewidth=2)
        axes[0, 2].plot(epochs, val_r2, 'r-', label='Val R²', linewidth=2)
        axes[0, 2].set_xlabel('Epoch')
        axes[0, 2].set_ylabel('R²')
        axes[0, 2].set_title('R² Score')
        axes[0, 2].legend()
        axes[0, 2].grid(True)
        
        # 4. RMSE曲线
        train_rmse = [m['RMSE'] for m in history['train_metrics']]
        val_rmse = [m['RMSE'] for m in history['val_metrics']]
        axes[1, 0].plot(epochs, train_rmse, 'b-', label='Train RMSE', linewidth=2)
        axes[1, 0].plot(epochs, val_rmse, 'r-', label='Val RMSE', linewidth=2)
        axes[1, 0].set_xlabel('Epoch')
        axes[1, 0].set_ylabel('RMSE')
        axes[1, 0].set_title('Root Mean Squared Error')
        axes[1, 0].legend()
        axes[1, 0].grid(True)
        
        # 5. 学习率曲线
        axes[1, 1].plot(epochs, history['learning_rates'], 'g-', linewidth=2)
        axes[1, 1].set_xlabel('Epoch')
        axes[1, 1].set_ylabel('Learning Rate')
        axes[1, 1].set_title('Learning Rate Schedule')
        axes[1, 1].set_yscale('log')
        axes[1, 1].grid(True)
        
        # 6. Accuracy曲线
        train_acc = [m['Accuracy@10%'] for m in history['train_metrics']]
        val_acc = [m['Accuracy@10%'] for m in history['val_metrics']]
        axes[1, 2].plot(epochs, train_acc, 'b-', label='Train Acc@10%', linewidth=2)
        axes[1, 2].plot(epochs, val_acc, 'r-', label='Val Acc@10%', linewidth=2)
        axes[1, 2].set_xlabel('Epoch')
        axes[1, 2].set_ylabel('Accuracy (%)')
        axes[1, 2].set_title('Accuracy within 10% Error')
        axes[1, 2].legend()
        axes[1, 2].grid(True)
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Training history plot saved to: {save_path}")
    
    def plot_predictions(self, y_true, y_pred, save_name='predictions.png'):
        """
        绘制预测结果对比
        
        参数:
            y_true: 真实值
            y_pred: 预测值
            save_name: 保存文件名
        """
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        fig.suptitle('Prediction Analysis', fontsize=16, fontweight='bold')
        
        y_true = np.array(y_true).flatten()
        y_pred = np.array(y_pred).flatten()
        errors = y_pred - y_true
        
        # 1. 预测值 vs 真实值散点图
        axes[0, 0].scatter(y_true, y_pred, alpha=0.5, s=20)
        
        # 添加对角线（完美预测线）
        min_val = min(y_true.min(), y_pred.min())
        max_val = max(y_true.max(), y_pred.max())
        axes[0, 0].plot([min_val, max_val], [min_val, max_val], 'r--', linewidth=2, label='Perfect Prediction')
        
        axes[0, 0].set_xlabel('True Values')
        axes[0, 0].set_ylabel('Predicted Values')
        axes[0, 0].set_title('Predictions vs True Values')
        axes[0, 0].legend()
        axes[0, 0].grid(True, alpha=0.3)
        
        # 添加R²和MAE信息
        from sklearn.metrics import r2_score, mean_absolute_error
        r2 = r2_score(y_true, y_pred)
        mae = mean_absolute_error(y_true, y_pred)
        axes[0, 0].text(0.05, 0.95, f'R² = {r2:.4f}\nMAE = {mae:.4f}',
                       transform=axes[0, 0].transAxes, fontsize=12,
                       verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        # 2. 误差分布直方图
        axes[0, 1].hist(errors, bins=50, edgecolor='black', alpha=0.7)
        axes[0, 1].axvline(x=0, color='r', linestyle='--', linewidth=2, label='Zero Error')
        axes[0, 1].set_xlabel('Prediction Error')
        axes[0, 1].set_ylabel('Frequency')
        axes[0, 1].set_title('Error Distribution')
        axes[0, 1].legend()
        axes[0, 1].grid(True, alpha=0.3)
        
        # 添加统计信息
        axes[0, 1].text(0.05, 0.95, 
                       f'Mean: {errors.mean():.4f}\nStd: {errors.std():.4f}',
                       transform=axes[0, 1].transAxes, fontsize=12,
                       verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        # 3. 残差图
        axes[1, 0].scatter(y_pred, errors, alpha=0.5, s=20)
        axes[1, 0].axhline(y=0, color='r', linestyle='--', linewidth=2)
        axes[1, 0].set_xlabel('Predicted Values')
        axes[1, 0].set_ylabel('Residuals')
        axes[1, 0].set_title('Residual Plot')
        axes[1, 0].grid(True, alpha=0.3)
        
        # 4. 相对误差分布
        relative_errors = np.abs(errors) / (y_true + 1e-8) * 100
        axes[1, 1].hist(relative_errors, bins=50, edgecolor='black', alpha=0.7)
        axes[1, 1].axvline(x=10, color='g', linestyle='--', linewidth=2, label='10% Error')
        axes[1, 1].axvline(x=20, color='orange', linestyle='--', linewidth=2, label='20% Error')
        axes[1, 1].set_xlabel('Relative Error (%)')
        axes[1, 1].set_ylabel('Frequency')
        axes[1, 1].set_title('Relative Error Distribution')
        axes[1, 1].legend()
        axes[1, 1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Prediction plot saved to: {save_path}")
    
    def plot_metrics_comparison(self, history, save_name='metrics_comparison.png'):
        """
        绘制多指标对比雷达图
        
        参数:
            history: 训练历史
            save_name: 保存文件名
        """
        # 获取最后一个epoch的指标
        final_train_metrics = history['train_metrics'][-1]
        final_val_metrics = history['val_metrics'][-1]
        
        # 选择关键指标（归一化到0-1）
        metrics_names = ['R²', 'Acc@10%', 'Acc@20%']
        
        train_values = [
            final_train_metrics['R²'],
            final_train_metrics['Accuracy@10%'] / 100,
            final_train_metrics['Accuracy@20%'] / 100
        ]
        
        val_values = [
            final_val_metrics['R²'],
            final_val_metrics['Accuracy@10%'] / 100,
            final_val_metrics['Accuracy@20%'] / 100
        ]
        
        # 雷达图
        angles = np.linspace(0, 2 * np.pi, len(metrics_names), endpoint=False).tolist()
        train_values += train_values[:1]
        val_values += val_values[:1]
        angles += angles[:1]
        
        fig, ax = plt.subplots(figsize=(10, 10), subplot_kw=dict(projection='polar'))
        
        ax.plot(angles, train_values, 'o-', linewidth=2, label='Train', color='blue')
        ax.fill(angles, train_values, alpha=0.25, color='blue')
        
        ax.plot(angles, val_values, 'o-', linewidth=2, label='Validation', color='red')
        ax.fill(angles, val_values, alpha=0.25, color='red')
        
        ax.set_xticks(angles[:-1])
        ax.set_xticklabels(metrics_names, size=12)
        ax.set_ylim(0, 1)
        ax.set_title('Model Performance Metrics', size=16, fontweight='bold', pad=20)
        ax.legend(loc='upper right', bbox_to_anchor=(1.3, 1.1))
        ax.grid(True)
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Metrics comparison plot saved to: {save_path}")
    
    def create_summary_report(self, history, best_epoch, save_name='summary_report.png'):
        """
        创建训练总结报告
        
        参数:
            history: 训练历史
            best_epoch: 最佳epoch
            save_name: 保存文件名
        """
        fig = plt.figure(figsize=(16, 10))
        fig.suptitle('Training Summary Report', fontsize=18, fontweight='bold')
        
        # 创建网格布局
        gs = fig.add_gridspec(3, 3, hspace=0.3, wspace=0.3)
        
        epochs = range(1, len(history['train_loss']) + 1)
        best_metrics = history['val_metrics'][best_epoch - 1]
        
        # 1. Loss曲线（大图）
        ax1 = fig.add_subplot(gs[0, :2])
        ax1.plot(epochs, history['train_loss'], 'b-', label='Train', linewidth=2, alpha=0.7)
        ax1.plot(epochs, history['val_loss'], 'r-', label='Validation', linewidth=2, alpha=0.7)
        ax1.axvline(x=best_epoch, color='g', linestyle='--', linewidth=2, label=f'Best Epoch ({best_epoch})')
        ax1.set_xlabel('Epoch', fontsize=12)
        ax1.set_ylabel('Loss', fontsize=12)
        ax1.set_title('Loss Curves', fontsize=14, fontweight='bold')
        ax1.legend(fontsize=10)
        ax1.grid(True, alpha=0.3)
        
        # 2. 最佳指标表格
        ax2 = fig.add_subplot(gs[0, 2])
        ax2.axis('off')
        
        metrics_text = f"""
        Best Epoch: {best_epoch}
        
        Validation Metrics:
        ━━━━━━━━━━━━━━━━
        MAE:     {best_metrics['MAE']:.4f}
        RMSE:    {best_metrics['RMSE']:.4f}
        R²:      {best_metrics['R²']:.4f}
        MAPE:    {best_metrics['MAPE']:.2f}%
        
        Acc@10%: {best_metrics['Accuracy@10%']:.2f}%
        Acc@20%: {best_metrics['Accuracy@20%']:.2f}%
        """
        
        ax2.text(0.1, 0.5, metrics_text, fontsize=11, verticalalignment='center',
                family='monospace', bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.5))
        
        # 3-5. 关键指标曲线
        metrics_to_plot = [
            ('MAE', 'Mean Absolute Error'),
            ('R²', 'R² Score'),
            ('Accuracy@10%', 'Accuracy@10%')
        ]
        
        for idx, (metric_key, metric_title) in enumerate(metrics_to_plot):
            ax = fig.add_subplot(gs[1, idx])
            train_vals = [m[metric_key] for m in history['train_metrics']]
            val_vals = [m[metric_key] for m in history['val_metrics']]
            
            ax.plot(epochs, train_vals, 'b-', label='Train', linewidth=2, alpha=0.7)
            ax.plot(epochs, val_vals, 'r-', label='Val', linewidth=2, alpha=0.7)
            ax.axvline(x=best_epoch, color='g', linestyle='--', linewidth=1, alpha=0.5)
            ax.set_xlabel('Epoch', fontsize=10)
            ax.set_ylabel(metric_key, fontsize=10)
            ax.set_title(metric_title, fontsize=12)
            ax.legend(fontsize=9)
            ax.grid(True, alpha=0.3)
        
        # 6. 学习率曲线
        ax6 = fig.add_subplot(gs[2, 0])
        ax6.plot(epochs, history['learning_rates'], 'purple', linewidth=2)
        ax6.set_xlabel('Epoch', fontsize=10)
        ax6.set_ylabel('Learning Rate', fontsize=10)
        ax6.set_title('Learning Rate Schedule', fontsize=12)
        ax6.set_yscale('log')
        ax6.grid(True, alpha=0.3)
        
        # 7. 训练vs验证Loss对比
        ax7 = fig.add_subplot(gs[2, 1])
        final_train_loss = history['train_loss'][-1]
        final_val_loss = history['val_loss'][-1]
        ax7.bar(['Train', 'Validation'], [final_train_loss, final_val_loss], 
               color=['blue', 'red'], alpha=0.7)
        ax7.set_ylabel('Final Loss', fontsize=10)
        ax7.set_title('Final Loss Comparison', fontsize=12)
        ax7.grid(True, alpha=0.3, axis='y')
        
        # 8. 训练统计
        ax8 = fig.add_subplot(gs[2, 2])
        ax8.axis('off')
        
        total_epochs = len(history['train_loss'])
        final_lr = history['learning_rates'][-1]
        
        stats_text = f"""
        Training Statistics:
        ━━━━━━━━━━━━━━━━
        Total Epochs: {total_epochs}
        Best Epoch:   {best_epoch}
        Final LR:     {final_lr:.2e}
        
        Improvement:
        Val Loss:     {history['val_loss'][0]:.4f} → {history['val_loss'][-1]:.4f}
        Val R²:       {history['val_metrics'][0]['R²']:.4f} → {history['val_metrics'][-1]['R²']:.4f}
        """
        
        ax8.text(0.1, 0.5, stats_text, fontsize=10, verticalalignment='center',
                family='monospace', bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.5))
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Summary report saved to: {save_path}")
        
    def plot_loss_components(self, history, save_name='loss_components.png'):
        """
        损失分解图 - 显示训练集和验证集损失的差异
        """
        fig, axes = plt.subplots(1, 2, figsize=(16, 6))
        
        epochs = range(1, len(history['train_loss']) + 1)
        train_loss = np.array(history['train_loss'])
        val_loss = np.array(history['val_loss'])
        gap = val_loss - train_loss
        
        # 左图：损失对比
        axes[0].plot(epochs, train_loss, 'b-', label='Training Loss', linewidth=2)
        axes[0].plot(epochs, val_loss, 'r-', label='Validation Loss', linewidth=2)
        axes[0].fill_between(epochs, train_loss, val_loss, alpha=0.3, color='gray', label='Generalization Gap')
        axes[0].set_xlabel('Epoch', fontsize=12)
        axes[0].set_ylabel('Loss', fontsize=12)
        axes[0].set_title('Training vs Validation Loss', fontsize=14, fontweight='bold')
        axes[0].legend(fontsize=11)
        axes[0].grid(True, alpha=0.3)
        
        # 右图：泛化差距
        axes[1].plot(epochs, gap, 'purple', linewidth=2)
        axes[1].axhline(y=0, color='black', linestyle='--', linewidth=1)
        axes[1].fill_between(epochs, 0, gap, alpha=0.3, color='purple')
        axes[1].set_xlabel('Epoch', fontsize=12)
        axes[1].set_ylabel('Generalization Gap (Val - Train)', fontsize=12)
        axes[1].set_title('Model Generalization Analysis', fontsize=14, fontweight='bold')
        axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Loss components plot saved to: {save_path}")


    def plot_error_analysis(self, y_true, y_pred, save_name='error_analysis.png'):
        """
        详细误差分析图
        """
        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
        fig.suptitle('Comprehensive Error Analysis', fontsize=16, fontweight='bold')
        
        y_true = np.array(y_true).flatten()
        y_pred = np.array(y_pred).flatten()
        errors = y_pred - y_true
        abs_errors = np.abs(errors)
        relative_errors = abs_errors / (y_true + 1e-8) * 100
        
        # 1. 按真实值分段的误差箱线图
        n_bins = 5
        bins = np.percentile(y_true, np.linspace(0, 100, n_bins + 1))
        bin_indices = np.digitize(y_true, bins[1:-1])
        
        bin_errors = [abs_errors[bin_indices == i] for i in range(n_bins)]
        axes[0, 0].boxplot(bin_errors, labels=[f'Q{i+1}' for i in range(n_bins)])
        axes[0, 0].set_xlabel('True Value Quantile', fontsize=11)
        axes[0, 0].set_ylabel('Absolute Error', fontsize=11)
        axes[0, 0].set_title('Error Distribution by Value Range', fontsize=12)
        axes[0, 0].grid(True, alpha=0.3, axis='y')
        
        # 2. Q-Q图（检查误差正态性）
        from scipy import stats
        stats.probplot(errors, dist="norm", plot=axes[0, 1])
        axes[0, 1].set_title('Q-Q Plot (Normality Check)', fontsize=12)
        axes[0, 1].grid(True, alpha=0.3)
        
        # 3. 累积误差分布
        sorted_rel_errors = np.sort(relative_errors)
        cumulative = np.arange(1, len(sorted_rel_errors) + 1) / len(sorted_rel_errors) * 100
        axes[0, 2].plot(sorted_rel_errors, cumulative, linewidth=2)
        axes[0, 2].axvline(x=10, color='g', linestyle='--', label='10% Error')
        axes[0, 2].axvline(x=20, color='orange', linestyle='--', label='20% Error')
        axes[0, 2].set_xlabel('Relative Error (%)', fontsize=11)
        axes[0, 2].set_ylabel('Cumulative Percentage (%)', fontsize=11)
        axes[0, 2].set_title('Cumulative Error Distribution', fontsize=12)
        axes[0, 2].legend()
        axes[0, 2].grid(True, alpha=0.3)
        
        # 4. 误差随预测值的变化（Bland-Altman图）
        mean_vals = (y_true + y_pred) / 2
        axes[1, 0].scatter(mean_vals, errors, alpha=0.5, s=20)
        mean_error = errors.mean()
        std_error = errors.std()
        axes[1, 0].axhline(mean_error, color='red', linestyle='-', linewidth=2, label='Mean Error')
        axes[1, 0].axhline(mean_error + 1.96 * std_error, color='red', linestyle='--', linewidth=1.5, label='±1.96 SD')
        axes[1, 0].axhline(mean_error - 1.96 * std_error, color='red', linestyle='--', linewidth=1.5)
        axes[1, 0].set_xlabel('Mean of True and Predicted', fontsize=11)
        axes[1, 0].set_ylabel('Difference (Predicted - True)', fontsize=11)
        axes[1, 0].set_title('Bland-Altman Plot', fontsize=12)
        axes[1, 0].legend()
        axes[1, 0].grid(True, alpha=0.3)
        
        # 5. 误差热力图（2D密度图）
        from scipy.stats import gaussian_kde
        xy = np.vstack([y_true, y_pred])
        z = gaussian_kde(xy)(xy)
        scatter = axes[1, 1].scatter(y_true, y_pred, c=z, s=20, cmap='YlOrRd', alpha=0.6)
        plt.colorbar(scatter, ax=axes[1, 1], label='Density')
        
        min_val = min(y_true.min(), y_pred.min())
        max_val = max(y_true.max(), y_pred.max())
        axes[1, 1].plot([min_val, max_val], [min_val, max_val], 'b--', linewidth=2)
        axes[1, 1].set_xlabel('True Values', fontsize=11)
        axes[1, 1].set_ylabel('Predicted Values', fontsize=11)
        axes[1, 1].set_title('Prediction Density Map', fontsize=12)
        
        # 6. Top-K最大误差样本
        top_k = 10
        worst_indices = np.argsort(abs_errors)[-top_k:][::-1]
        
        x_pos = np.arange(top_k)
        axes[1, 2].barh(x_pos, abs_errors[worst_indices], color='red', alpha=0.7)
        axes[1, 2].set_yticks(x_pos)
        axes[1, 2].set_yticklabels([f'#{i+1}' for i in range(top_k)])
        axes[1, 2].set_xlabel('Absolute Error', fontsize=11)
        axes[1, 2].set_ylabel('Sample Rank', fontsize=11)
        axes[1, 2].set_title(f'Top-{top_k} Largest Errors', fontsize=12)
        axes[1, 2].grid(True, alpha=0.3, axis='x')
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Error analysis plot saved to: {save_path}")


    def plot_metrics_evolution(self, history, save_name='metrics_evolution.png'):
        """
        多指标演化趋势图
        """
        fig, axes = plt.subplots(3, 4, figsize=(20, 15))
        fig.suptitle('Comprehensive Metrics Evolution', fontsize=18, fontweight='bold')
        
        epochs = range(1, len(history['train_loss']) + 1)
        
        # 定义所有要绘制的指标
        metrics_list = [
            'MAE', 'RMSE', 'R²', 'MAPE',
            'MedAE', 'MaxError', 'Pearson', 'Spearman',
            'Accuracy@10%', 'Accuracy@20%', 'Std_Error', 'MSLE'
        ]
        
        for idx, metric in enumerate(metrics_list):
            row = idx // 4
            col = idx % 4
            ax = axes[row, col]
            
            train_vals = [m[metric] for m in history['train_metrics']]
            val_vals = [m[metric] for m in history['val_metrics']]
            
            ax.plot(epochs, train_vals, 'b-', label='Train', linewidth=2, alpha=0.7)
            ax.plot(epochs, val_vals, 'r-', label='Val', linewidth=2, alpha=0.7)
            ax.set_xlabel('Epoch', fontsize=10)
            ax.set_ylabel(metric, fontsize=10)
            ax.set_title(metric, fontsize=12, fontweight='bold')
            ax.legend(fontsize=9)
            ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Metrics evolution plot saved to: {save_path}")


    def plot_learning_curve(self, history, save_name='learning_curve.png'):
        """
        学习曲线（训练集大小 vs 性能）
        """
        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
        fig.suptitle('Learning Curves Analysis', fontsize=16, fontweight='bold')
        
        epochs = range(1, len(history['train_loss']) + 1)
        
        # 1. Loss收敛曲线（对数尺度）
        axes[0].semilogy(epochs, history['train_loss'], 'b-', label='Train Loss', linewidth=2)
        axes[0].semilogy(epochs, history['val_loss'], 'r-', label='Val Loss', linewidth=2)
        axes[0].set_xlabel('Epoch', fontsize=12)
        axes[0].set_ylabel('Loss (log scale)', fontsize=12)
        axes[0].set_title('Loss Convergence (Log Scale)', fontsize=14)
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)
        
        # 2. 训练效率分析
        train_loss_improvement = np.array(history['train_loss'])
        train_loss_rate = -np.diff(train_loss_improvement, prepend=train_loss_improvement[0])
        
        axes[1].plot(epochs, train_loss_rate, 'g-', linewidth=2)
        axes[1].axhline(y=0, color='black', linestyle='--', linewidth=1)
        axes[1].set_xlabel('Epoch', fontsize=12)
        axes[1].set_ylabel('Loss Improvement Rate', fontsize=12)
        axes[1].set_title('Training Efficiency', fontsize=14)
        axes[1].grid(True, alpha=0.3)
        
        # 3. 过拟合风险分析
        train_r2 = np.array([m['R²'] for m in history['train_metrics']])
        val_r2 = np.array([m['R²'] for m in history['val_metrics']])
        overfitting_risk = train_r2 - val_r2
        
        axes[2].plot(epochs, overfitting_risk, 'orange', linewidth=2)
        axes[2].axhline(y=0, color='green', linestyle='--', linewidth=2, label='No Overfitting')
        axes[2].axhline(y=0.1, color='orange', linestyle='--', linewidth=1, label='Mild Overfitting')
        axes[2].axhline(y=0.2, color='red', linestyle='--', linewidth=1, label='Severe Overfitting')
        axes[2].set_xlabel('Epoch', fontsize=12)
        axes[2].set_ylabel('R² Gap (Train - Val)', fontsize=12)
        axes[2].set_title('Overfitting Risk', fontsize=14)
        axes[2].legend()
        axes[2].grid(True, alpha=0.3)
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Learning curve plot saved to: {save_path}")


    def plot_prediction_intervals(self, y_true, y_pred, save_name='prediction_intervals.png'):
        """
        预测区间和置信度分析
        """
        fig, axes = plt.subplots(1, 2, figsize=(16, 6))
        fig.suptitle('Prediction Confidence Analysis', fontsize=16, fontweight='bold')
        
        y_true = np.array(y_true).flatten()
        y_pred = np.array(y_pred).flatten()
        errors = y_pred - y_true
        
        # 排序以便绘制
        sorted_indices = np.argsort(y_true)
        y_true_sorted = y_true[sorted_indices]
        y_pred_sorted = y_pred[sorted_indices]
        errors_sorted = errors[sorted_indices]
        
        # 计算置信区间（使用局部标准差）
        window_size = max(len(y_true) // 20, 10)
        std_errors = []
        for i in range(len(y_true)):
            start = max(0, i - window_size // 2)
            end = min(len(y_true), i + window_size // 2)
            std_errors.append(np.std(errors_sorted[start:end]))
        std_errors = np.array(std_errors)
        
        # 1. 预测区间图
        axes[0].plot(y_true_sorted, y_true_sorted, 'g--', linewidth=2, label='Perfect Prediction', alpha=0.7)
        axes[0].plot(y_true_sorted, y_pred_sorted, 'b-', linewidth=2, label='Predictions')
        axes[0].fill_between(y_true_sorted, 
                            y_pred_sorted - 1.96 * std_errors,
                            y_pred_sorted + 1.96 * std_errors,
                            alpha=0.3, color='blue', label='95% Confidence Interval')
        axes[0].set_xlabel('True Values', fontsize=12)
        axes[0].set_ylabel('Predicted Values', fontsize=12)
        axes[0].set_title('Predictions with Confidence Intervals', fontsize=14)
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)
        
        # 2. 误差标准差分布
        axes[1].plot(y_true_sorted, std_errors, 'r-', linewidth=2)
        axes[1].fill_between(y_true_sorted, 0, std_errors, alpha=0.3, color='red')
        axes[1].set_xlabel('True Values', fontsize=12)
        axes[1].set_ylabel('Prediction Uncertainty (Std)', fontsize=12)
        axes[1].set_title('Prediction Uncertainty Distribution', fontsize=14)
        axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Prediction intervals plot saved to: {save_path}")


    def create_comparison_table(self, history, best_epoch, save_name='metrics_table.png'):
        """
        创建详细的指标对比表格图
        """
        fig, ax = plt.subplots(figsize=(14, 10))
        ax.axis('off')
        
        # 获取最佳epoch和最终epoch的指标
        best_train_metrics = history['train_metrics'][best_epoch - 1]
        best_val_metrics = history['val_metrics'][best_epoch - 1]
        final_train_metrics = history['train_metrics'][-1]
        final_val_metrics = history['val_metrics'][-1]
        
        # 准备表格数据
        metrics_names = ['MAE', 'RMSE', 'R²', 'MAPE (%)', 'MedAE', 'MaxError', 
                        'Pearson', 'Spearman', 'Acc@10% (%)', 'Acc@20% (%)']
        
        table_data = []
        for metric in metrics_names:
            metric_key = metric.replace(' (%)', '')
            if metric_key in best_train_metrics:
                table_data.append([
                    metric,
                    f"{best_train_metrics[metric_key]:.4f}",
                    f"{best_val_metrics[metric_key]:.4f}",
                    f"{final_train_metrics[metric_key]:.4f}",
                    f"{final_val_metrics[metric_key]:.4f}"
                ])
        
        # 创建表格
        table = ax.table(cellText=table_data,
                        colLabels=['Metric', f'Train\n(Epoch {best_epoch})', f'Val\n(Epoch {best_epoch})', 
                                f'Train\n(Final)', f'Val\n(Final)'],
                        cellLoc='center',
                        loc='center',
                        colWidths=[0.2, 0.2, 0.2, 0.2, 0.2])
        
        table.auto_set_font_size(False)
        table.set_fontsize(11)
        table.scale(1, 2.5)
        
        # 设置表头样式
        for i in range(5):
            table[(0, i)].set_facecolor('#4CAF50')
            table[(0, i)].set_text_props(weight='bold', color='white')
        
        # 交替行颜色
        for i in range(1, len(table_data) + 1):
            for j in range(5):
                if i % 2 == 0:
                    table[(i, j)].set_facecolor('#f0f0f0')
        
        plt.title('Detailed Metrics Comparison Table', fontsize=16, fontweight='bold', pad=20)
        
        plt.tight_layout()
        save_path = self.save_dir / save_name
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Metrics table saved to: {save_path}")